package common

import (
	"sync"
)

// ============================================================================
// 缓冲池 - 减少内存分配开销
// ============================================================================
//
// 问题：
// - 每个数据包都分配新内存 (make([]byte, size))
// - 频繁的GC导致性能下降
//
// 解决方案：
// - 使用 sync.Pool 复用缓冲区
// - 分层缓冲池（小/中/大）
// - 减少内存分配和GC压力
//
// ============================================================================

const (
	SmallBufferSize  = 4 * 1024    // 4KB - 小数据包
	MediumBufferSize = 64 * 1024   // 64KB - 中等数据包
	LargeBufferSize  = 1024 * 1024 // 1MB - 大数据包
)

var (
	// 小缓冲池 (4KB) - 用于大多数请求/响应
	smallBufferPool = sync.Pool{
		New: func() interface{} {
			b := make([]byte, SmallBufferSize)
			return &b
		},
	}

	// 中缓冲池 (64KB) - 用于较大的响应
	mediumBufferPool = sync.Pool{
		New: func() interface{} {
			b := make([]byte, MediumBufferSize)
			return &b
		},
	}

	// 大缓冲池 (1MB) - 用于大文件传输
	largeBufferPool = sync.Pool{
		New: func() interface{} {
			b := make([]byte, LargeBufferSize)
			return &b
		},
	}
)

// GetBuffer 从缓冲池获取合适大小的缓冲区
func GetBuffer(size int) []byte {
	if size <= SmallBufferSize {
		bufPtr := smallBufferPool.Get().(*[]byte)
		return (*bufPtr)[:size] // 返回需要的长度
	} else if size <= MediumBufferSize {
		bufPtr := mediumBufferPool.Get().(*[]byte)
		return (*bufPtr)[:size]
	} else if size <= LargeBufferSize {
		bufPtr := largeBufferPool.Get().(*[]byte)
		return (*bufPtr)[:size]
	} else {
		// 超大缓冲区，直接分配（不使用池）
		return make([]byte, size)
	}
}

// PutBuffer 将缓冲区归还到池中
func PutBuffer(buf []byte) {
	capacity := cap(buf)
	if capacity == SmallBufferSize {
		bufPtr := &buf
		*bufPtr = (*bufPtr)[:SmallBufferSize] // 恢复完整长度
		smallBufferPool.Put(bufPtr)
	} else if capacity == MediumBufferSize {
		bufPtr := &buf
		*bufPtr = (*bufPtr)[:MediumBufferSize]
		mediumBufferPool.Put(bufPtr)
	} else if capacity == LargeBufferSize {
		bufPtr := &buf
		*bufPtr = (*bufPtr)[:LargeBufferSize]
		largeBufferPool.Put(bufPtr)
	}
	// 超大缓冲区不归还，让GC处理
}

// GetSmallBuffer 快速获取小缓冲区（4KB）
func GetSmallBuffer() []byte {
	bufPtr := smallBufferPool.Get().(*[]byte)
	return *bufPtr
}

// PutSmallBuffer 归还小缓冲区
func PutSmallBuffer(buf []byte) {
	if cap(buf) != SmallBufferSize {
		return // 大小不匹配，不归还
	}
	bufPtr := &buf
	smallBufferPool.Put(bufPtr)
}

// GetMediumBuffer 快速获取中缓冲区（64KB）
func GetMediumBuffer() []byte {
	bufPtr := mediumBufferPool.Get().(*[]byte)
	return *bufPtr
}

// PutMediumBuffer 归还中缓冲区
func PutMediumBuffer(buf []byte) {
	if cap(buf) != MediumBufferSize {
		return
	}
	bufPtr := &buf
	mediumBufferPool.Put(bufPtr)
}

// GetLargeBuffer 快速获取大缓冲区（1MB）
func GetLargeBuffer() []byte {
	bufPtr := largeBufferPool.Get().(*[]byte)
	return *bufPtr
}

// PutLargeBuffer 归还大缓冲区
func PutLargeBuffer(buf []byte) {
	if cap(buf) != LargeBufferSize {
		return
	}
	bufPtr := &buf
	largeBufferPool.Put(bufPtr)
}
